///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file davidson.h
///
/// \brief   header file for the Davidson solver
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef PCPS_DAVIDSON_HEADER
#define PCPS_DAVIDSON_HEADER

#include <src/pcps.h>
#include <src/matrix_actor.h>
#include <src/mpi_interface.h>
#include <src/blas_lapack.h>

namespace pcps {

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  Solves a generalized hermitian eigenvalue problem  ( H x = E S x )  for the
  ///         (eigenvalue,eigenvector) pair (E,x) with the lowest eigenvalue using
  ///         the iterative Davidson solver.
  ///         
  ///         The solver constructs the matrix H in a subspace spanned by the vectors
  ///         (v_0, v_1, ...) and solves for the lowest eigenpair in this subspace.
  ///
  ///         The subpace is expanded iteratively by adding the vector
  ///
  ///         v_(i+1) = M ( H - E_i S ) x_i
  ///
  ///         where (E_i,x_i) is the current approximation to the desired eigenpair and M
  ///         is a preconditioning matrix.
  ///
  ///         Typically, this preconditioning matrix is chosen to be
  ///
  ///         ( diag(H) - E_i diag(S) )^(-1)
  ///
  ///         where diag(A) means the matrix with diagonal elements from A and zero off-diagonals.
  ///
  ///         In general, however, the matrix M may be chosen to be any matrix that is a good
  ///         approximation to the inverse of the matrix pencil ( H - E_i S ).
  ///
  /// \param[in]      n        the dimension of the matrices and vectors
  /// \param[in]      m        the maximum number of iterations
  /// \param[in]      thresh   threshold for convergence
  /// \param[in,out]  E        on entry, a guess for the eigenvalue.
  ///                          on exit, the subspace-approximated eigenvalue.
  /// \param[in,out]  x        on entry, a guess for the eigenvector.
  ///                          on exit, the subspace-approximated eigenvector.
  /// \param[in,out]  H        an object that multiplies by the Hermitian matrix H
  /// \param[in,out]  S        an object that multiplies by the Hermitian, positive definite matrix S
  /// \param[in,out]  M        an object that multiplies by the matrix M
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////

  template<class T> inline void davidson(const int n,
                                         const int m,
                                         const double thresh,
                                         double & E,
                                         T * const x,
                                         pcps::MatrixActor<T> * const H,
                                         pcps::MatrixActor<T> * const S,
                                         pcps::MatrixActor<T> * const M) {

    // make sure the dimension is sane
    if (n <= 0)
      throw pcps::Exception("invalid matrix dimension of \"%i\" in pcps::davidson") % n;

    // make sure the number of iterations is sane
    if (m <= 0)
      throw pcps::Exception("\"%i\" is not a valid number of iterations in pcps::davidson") % m;

    // initialize the residual vector
    std::vector<T> w(n);

    // create a vector to hold the subspace basis vectors (v_0, v_1, ...)
    std::vector<std::vector<T> > v(1);
    v[0].resize(n);

    // set the initial guess as the first basis vector
    pcps::xcopy(n, x, 1, &v[0][0], 1);

    // initialize the subspace matrix
    std::vector<T> subH;

    // create vectors of vectors to hold matrix-vector products
    std::vector<std::vector<T> > Hv, Sv;

    // initialize the residual norm
    double rnorm = 1.0e100;

    // iteratively expand the subspace matrices and basis vectors
    int iter = 0;
    bool converged = false;
    for ( ; iter < m && !converged; iter++) {

      // get the new subspace dimension
      const int p = iter + 1;

      // orthogonalize the new vector against the previous vectors
      for (int i = 0; i < iter; i++)
        pcps::xaxpy(n, -pcps::xdotc(n, &Sv[i][0], 1, &v[iter][0], 1), &v[i][0], 1, &v[iter][0], 1);

      // normalize the new basis vector
      S->operate(n, &v[iter][0], x);
      const T vnorm = std::sqrt(pcps::xdotc(n, &v[iter][0], 1, x, 1));
      pcps::xscal(n, pcps::unity<T>() / vnorm, &v[iter][0], 1);

      // place the new  S v  product in the Sv vector
      Sv.resize(p);
      Sv[iter].assign(n, pcps::zero<T>());
      pcps::xaxpy(n, pcps::unity<T>() / vnorm, x, 1, &Sv[iter][0], 1);

      // create a new matrix for the expanded subspace
      std::vector<T> newH(p*p, pcps::zero<T>());

      // copy the old matrix into the new matrix's upper left block
      for (int i = 0; i < iter; i++)
        for (int j = 0; j < iter; j++)
          newH[ p*i + j ] = subH[ iter*i + j ];

      // compute and store the new matrix vector product
      Hv.resize(p);
      Hv[iter].resize(n);
      H->operate(n, &v[iter][0], &Hv[iter][0]);

      // compute the new row and column of the subspace matrix
      for (int i = 0; i < p; i++) {
        newH[ p*i + iter ] = pcps::xdotc(n, &v[i][0], 1, &Hv[iter][0], 1);
        newH[ p*iter + i ] = pcps::conj(newH[ p*i + iter ]);
      }

      // store the new subspace matrix
      subH = newH;

      // compute the eigenvector and eigenvalue in the subspace
      std::vector<T> A(subH);
      for (int i = 0; i < p; i++)
        for (int j = 0; j < i; j++)
          std::swap(A[i*p+j], A[j*p+i]); // transpose for FORTRAN ordering
      std::vector<double> evals(p, 0.0);
      int lwork = 5*p;
      std::vector<T> work(lwork, 0.0);
      std::vector<double> rwork(lwork, 0.0);
      int info;
      pcps::xsyev('V', 'U', p, &A[0], p, &evals[0], &work[0], lwork, &rwork[0], info);

      // check that the eivenvalue solver was successful
      if (info != 0)
        throw pcps::Exception("pcps::xsyev failed with info = %i in pcps::davidson") % info;

      // extract the eigenvalue
      E = evals[0];

      // construct the new residual
      std::fill(w.begin(), w.end(), pcps::zero<T>());
      for (int i = 0; i < p; i++) {
        pcps::xaxpy(n, A[i], &Hv[i][0], 1, &w[0], 1);
        pcps::xaxpy(n, -E * A[i], &Sv[i][0], 1, &w[0], 1);
      }

      // compute the residual norm
      S->operate(n, &w[0], x);
      rnorm = std::abs(std::sqrt(pcps::xdotc(n, &w[0], 1, x, 1)));

      // compute the new subspace basis vector
      v.resize(p+1);
      v[p].resize(n);
      M->operate(n, &w[0], &v[p][0]);

      // construct the new eigenvector
      std::fill(x, x+n, pcps::zero<T>());
      for (int i = 0; i < p; i++)
        pcps::xaxpy(n, A[i], &v[i][0], 1, x, 1);

      // check if we are finished  (only the result on the root process is meaningful)
      converged = ( rnorm < thresh );
      pcps::mpi::bcast(converged);

      if (pcps::mpi::rank() == 0)
        std::cout << boost::format("davidson iteration %4i:          eval = %20.12f          resid = %20.12f")
                     % iter
                     % E
                     % rnorm
                  << std::endl;

    }

    // check if the method converged
    if (converged && pcps::mpi::rank() == 0)
      std::cout << boost::format("davidson solver converged in %10i iterations") % iter << std::endl << std::endl;
    else if (pcps::mpi::rank() == 0)
      std::cout << boost::format("davidson solver did not converge after %10i iterations") % iter << std::endl << std::endl;
      //throw pcps::Exception("pcps::davidson did not converge.  residual = %.2e") % std::abs(rnorm);

  }

  template<class T> void test_davidson(const int n);

} // end namespace pcps

#endif
